import os
from abc import ABC
from typing import Dict, Any, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import (
    BaseModel,
    Field,
    root_validator,
)
from loguru import logger

from ai_engine.common.ai_common import split_list
from ai_engine.core.llm.volcengine.ark_common import EmbModel, RawData
from ai_engine.core.llm.volcengine.ark_embeddings_service import ArkEmbeddingsService


class VolcEngineArkEmbeddings(BaseModel, Embeddings, ABC):
    """火山方舟向量模型"""
    # 火上方舟client
    client: Any
    # 模型/模型接入点
    model: str = "bge-m3"
    volc_ak: str | None = None,
    volc_sk: str | None = None,
    """What sampling temperature to use."""
    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
    chunk_size: int = 16
    request_timeout: Optional[int] = Field(30, alias="timeout")

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        try:
            if values["model"] is None or values["model"] == "":
                raise ValueError("must set model")
            ak = values["volc_ak"]
            sk = values["volc_sk"]
            if ak is None:
                ak = os.environ.get("VOLC_ACCESSKEY")
            if sk is None:
                sk = os.environ.get("VOLC_SECRETKEY")
            values["client"] = ArkEmbeddingsService(ak=ak, sk=sk)
        except ImportError:
            raise ImportError(
                "volcenginesdk package not found, please install it with "
                "`pip install  volcengine-python-sdk `"
            )
        return values

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed search docs."""
        if len(texts) > 100:
            logger.warning("embed_documents，texts>100")
        embedding_data = []
        # texts 一次最多支持100条文本生成向量
        chunked_list = split_list(texts, 100)
        for chunk in chunked_list:
            req_lst = []
            for text in chunk:
                req_lst.append(RawData("text", text))
            resp = self.client.create(EmbModel("bge-m3", params={"return_token_usage": True}), req_lst)
            logger.info("embedding token_usage: " + str(resp['token_usage']))
            embedding_data += resp['sentence_dense_embedding']
        return embedding_data

    def embed_query(self, text: str) -> List[float]:
        """Embed query text.
           Doubao-embedding 模型用于检索场景，query 建议必须添加instruction前缀保证检索效果。
        """
        # query_instruction = '为这个句子生成表示以用于检索相关文章：'
        resp = self.embed_documents([text])
        return resp[0]

# if __name__ == '__main__':
#
#     my_list=[{"categories":"e0091aa7803245638fd176bea4b77dc8","id":"19918f1567944a73a70990c870a58074",
#              "originalText":"图雅诺落地价会根据选购配置、个人情况等多种因素导致有所不同，不方便在平台上聊太细。优惠空间挺大的，
#              成为我们的粉丝还会有专门的到店礼、额外优惠。。现在全款跟分期都有不同的福利政策，不过具体首付分期这些会跟您本人的征信各方面有关，
#              可以跟顾问详细聊。。有置换补贴、官方补贴等多重政策福利，具体补贴方案可以叫顾问给您详细介绍。。质保都是跟着官方政策走的，这个不用担心的~。
#              保养建议直接找顾问哈，他会帮您安排的。。\n图雅诺2022款 2.8T手动吉象多功能加长轴中顶后双胎5/6/7/9座F2.8NS6B177L优惠空间挺大的，
#              成为我们的粉丝还会有专门的到店礼、额外优惠。\n图雅诺2023款 2.8T 小客豪华版手动短轴低顶后单胎7座营运优惠空间挺大的，成为我们的粉丝还会有专门的到店礼、
#              额外优惠。\n图雅诺2023款 2.8T 小客豪华版手动短轴中顶后单胎7座营运优惠空间挺大的，成为我们的粉丝还会有专门的到店礼、额外优惠。
#              \n图雅诺2023款 2.8T 小客豪华版手动短轴低顶后单胎7座优惠空间挺大的，成为我们的粉丝还会有专门的到店礼、额外优惠。
#              \n图雅诺2023款 2.8T 小客豪华版手动短轴中顶后单胎7座优惠空间挺大的，成为我们的粉丝还会有专门的到店礼、额外优惠。\
#              n图雅诺2023款 2.8T 小客豪华版手动长轴中顶后单胎9座营运优惠空间挺大的，成为我们的粉丝还会有专门的到店礼、额外优惠。
#              \n图雅诺2023款 2.8T 小客舒适版手动长轴中顶后单胎9座营运优惠空间挺大的，成为我们的粉丝还会有专门的到店礼、额外优惠。}]
#
#     embedding = VolcEngineArkEmbeddings(
#         model="ep-20240524024731-t82d8",
#         volc_ak="AKLTYTZmZjU0NjEwYmYzNDhlNGI2NTZiN2E4YjBlODhjM2M",
#         volc_sk="TURrME1EZGlPVGszWlRjeU5EUmlNV0l4WVRZeE9UVTBaREUxTURVM1pqYw=="
#     )
#     for doc in my_list:
#         vector_text=doc["originalText"]
#         try:
#             print(len(vector_text))
#             embedding.embed_documents([vector_text])
#         except Exception as e:
#             print("vector_text", vector_text)
#             print("异常",e)
